import base64
import json
import os
from pyexpat import model
import random
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import requests
from urllib3.exceptions import InsecureRequestWarning

# 忽略 SSL 警告（适用于自签名证书）
requests.packages.urllib3.disable_warnings(category=InsecureRequestWarning)

# ===== 配置区域 =====
API_URL = "http://172.16.1.1:8001/v1/chat/completions" 
API_KEY = "EMPTY"  # 如果没有启用 API Key，设为 EMPTY 或任意值
IMAGE_DIR = "/Users/joy/Desktop/images"  # 存放图片的文件夹路径
PROMPT = "分析图中场景并生成详细描述。"
CONCURRENCY_LEVELS = [1, 4, 8, 16, 24]  # 并发级别
Model = "Qwen2.5-VL-32B-Instruct"
Max_Tokens = 4096

# 将图片转为 base64 编码
def image_to_base64(file_path):
    with open(file_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode("utf-8")


# 构造一个带图片的消息内容
# def build_message(prompt, base64_image):
#     return {
#         "model": Model,
#         "messages": [
#             {
#                 "role": "user",
#                 "content": [
#                     {"type": "text", "text": prompt},
#                     {
#                         "type": "image_url",
#                         "image_url": {
#                             "url": f"data:image/jpeg;base64,{base64_image}"
#                         },
#                     },
#                 ],
#             }
#         ],
#         "max_tokens": Max_Tokens,
#     }

def build_message(prompt):
    return {
        "model": Model,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": "http://8.154.30.196:5001/uploads/202506271354/1751003672427-80c84e67-d399-4c92-ae11-67fa9cc7b970.JPG"
                        },
                    },
                ],
            }
        ],
        "max_tokens": Max_Tokens,
    }

# 加载所有图片路径并预处理为 base64（可选）
def load_all_images(image_dir):
    image_files = [f for f in os.listdir(image_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    if not image_files:
        raise FileNotFoundError(f"在目录 {image_dir} 中未找到任何图片文件")
    print(f"已加载 {len(image_files)} 张图片用于测试")
    return [os.path.join(image_dir, f) for f in image_files]


# 单次请求函数
def send_request(payload):
    start_time = time.time()
    try:
        response = requests.post(
            API_URL,
            headers={
                "Content-Type": "application/json",
                "Authorization": f"Bearer {API_KEY}",
            },
            json=payload,
            verify=False  # 忽略 SSL 证书验证（仅限测试）
        )
        latency = time.time() - start_time
        if response.status_code == 200:
            result = response.json()
            total_tokens = result["usage"]["completion_tokens"]
            tokens_per_second = total_tokens / latency if latency > 0 else 0
            return latency, True, total_tokens, tokens_per_second
        else:
            print(f"Error: {response.status_code}, {response.text}")
            return latency, False, 0, 0
    except Exception as e:
        print(f"Exception: {e}")
        return time.time() - start_time, False, 0, 0


# 性能测试主函数
def run_benchmark(concurrency, num_requests, image_paths):
    latencies = []
    tokens_list = []
    tokens_per_second_list = []
    successes = 0

    print(f"\n[并发数={concurrency}] 开始 {num_requests} 次请求...")

    with ThreadPoolExecutor(max_workers=concurrency) as executor:
        futures = []
        for _ in range(num_requests):
            # image_path = random.choice(image_paths)
            # base64_image = image_to_base64(image_path)
            # payload = build_message(PROMPT, base64_image)

            payload = build_message(PROMPT)
            futures.append(executor.submit(send_request, payload))

        for future in as_completed(futures):
            latency, success, tokens, tokens_per_second = future.result()
            latencies.append(latency)
            tokens_list.append(tokens)
            tokens_per_second_list.append(tokens_per_second)
            if success:
                successes += 1
    if successes == 0:
        print("全部失败，请检查问题。")
        return ""
    avg_latency = sum(latencies) / len(latencies)
    total_time = max(latencies)
    total_tokens = sum(tokens_list)
    token_speed = total_tokens / total_time if total_time > 0 else 0
    avg_token_speed = token_speed / successes
    tps = num_requests / total_time if total_time > 0 else 0

    print(
        f"[并发数={concurrency}] "
        f"平均延迟: {avg_latency:.2f}s | "
        f"成功数/总数: {successes}/{num_requests} | "
        f"TPS: {tps:.2f} | "
        f"总token/s: {token_speed:.2f} | "
        f"单路token/s: {avg_token_speed:.2f}"
    )

    return {
        "concurrency": concurrency,
        "total_requests": num_requests,
        "successes": successes,
        "avg_latency": avg_latency,
        "tps": tps,
        "token_speed": token_speed,
        "avg_token_speed": avg_token_speed
    }


# 主程序入口
if __name__ == "__main__":
    # 加载所有图片路径
    image_paths = load_all_images(IMAGE_DIR)

    results = []
    for level in CONCURRENCY_LEVELS:
        result = run_benchmark(level, level, image_paths)
        if result == "":
            break
        results.append(result)

    print("\n=== 性能测试结果汇总 ===")
    print("{:<8} {:<8} {:<10} {:<10} {:<12} {:<12}".format(
        "并发数", "成功率", "平均延迟(s)", "TPS", "总 token/s", "平均 token/s"))
    for res in results:
        success_rate = f"{res['successes']}/{res['total_requests']}"
        print("{:<10} {:<10} {:<13.2f} {:<10.2f} {:<12.2f} {:<12.2f}".format(
            res['concurrency'],
            success_rate,
            res['avg_latency'],
            res['tps'],
            res['token_speed'],
            res['avg_token_speed']  # 确保此字段存在，如需请修改run_benchmark函数
        ))